from threading import Thread

from flask import request, jsonify
from src.app.services.ner_identify import ner_identify

def identify():
    req_data = request.json
    schema = None
    if 'schema' in req_data:
        schema = req_data['schema']
    return ner_identify.identify(req_data['data'], schema)

def train_model():
    return ner_identify.train()


# 全局变量存储训练进度
training_status = {
    'progress': 0,
    'status': 'idle',
    'epoch': '0/0',
    'iter': '0/0'
}

def start_training():
    def training_thread():
        def progress_callback(data):
            global training_status
            training_status = {
                'progress': data['progress'],
                'status': data['status'],
                'epoch': data['epoch'],
                'iter': data['iter']
            }

        ner_identify.train_ner_model(progress_callback=progress_callback)

    # 检查是否已经在训练
    if training_status['status'] not in ['idle', 'completed']:
        return jsonify({'error': 'Training is already in progress'}), 400

    # 重置状态
    training_status['status'] = 'starting'

    # 在新线程中启动训练
    thread = Thread(target=training_thread)
    thread.start()

    return jsonify({
        'code': 200,
        'message': 'training started'
    })


def get_training_status():
    return jsonify({
        'code': 200,
        'data': training_status
    })
